import matplotlib.pyplot as plt
from cec17.cec2017.functions import all_functions
import numpy as np
import os
import copy

# 种群初始化
def initialization(pop_size, dim, ub, lb):
    return np.random.rand(pop_size, dim) * (ub - lb) + lb

# 返回指定的cec2017数据集中的函数
def fhd(x, func_num):
    return all_functions[func_num - 1](x)

def PSO(population, pop_size, Tmax, ub, lb, dim, func_num, w=0.5, c1=1.5, c2=1.5):
    V = np.random.rand(pop_size, dim)
    pbest = np.copy(population)  # Personal best position
    pbest_fitness = fhd(pbest, func_num)  # Personal best fitness
    gbest = pbest[np.argmin(pbest_fitness), :]  # Global best position
    gbest_fitness = np.min(pbest_fitness)  # Global best fitness
    conv_curve = []

    for t in range(Tmax):
        for i in range(pop_size):
            r1 = np.random.rand(dim)
            r2 = np.random.rand(dim)
            V[i] = w * V[i] + c1 * r1 * (pbest[i] - population[i]) + c2 * r2 * (gbest - population[i])
            population[i] = np.clip(population[i] + V[i], lb, ub)
            fitness = fhd(population[np.newaxis, i], func_num).item()
            if fitness < pbest_fitness[i]:
                pbest[i] = population[i]
                pbest_fitness[i] = fitness
                if fitness < gbest_fitness:
                    gbest = population[i]
                    gbest_fitness = fitness
        conv_curve.append(gbest_fitness)
    return gbest_fitness, gbest, np.array(conv_curve)

# CPO算法(冠状猪)
def CPO(population, pop_size, Tmax, ub, lb, dim, func_num):
    Gb_Fit = np.inf
    Gb_Sol = None
    Conv_curve = np.zeros(Tmax)
    fitness = fhd(population, func_num)
    Gb_Fit, index = np.min(fitness), np.argmin(fitness)
    Gb_Sol = population[index, :]
    Xp = np.copy(population)
    opt = 0
    t = 0

    while t < Tmax and Gb_Fit > opt:
        for i in range(len(population)):
            U1 = np.random.rand(dim) > 0.5
            rand_index1 = np.random.randint(len(population))
            rand_index2 = np.random.randint(len(population))

            if np.random.rand() < 0.5:
                y = (population[i, :] + population[rand_index1, :]) / 2
                population[i, :] = population[i, :] + np.random.randn(dim) * np.abs(2 * np.random.rand() * Gb_Sol - y)
            else:
                Yt = 2 * np.random.rand() * (1 - t / Tmax) ** (t / Tmax + 1e-6)
                U2 = np.random.rand(dim) < 0.5
                S = np.random.rand() * U2
                if np.random.rand() < 0.8:
                    St = np.exp(-np.abs(fitness[i] / (np.sum(fitness) + np.finfo(float).eps)))
                    S = S * Yt * St
                    population[i, :] = (1 - U1) * population[i, :] + U1 * (
                            population[rand_index1, :] + St * (population[rand_index2, :] - population[rand_index1, :]) - S)
                else:
                    Mt = np.exp(-np.abs(fitness[i] / (np.sum(fitness) + np.finfo(float).eps)))
                    Vtp = population[rand_index1, :]
                    Ft = np.random.rand(dim) * (Mt * (-population[i, :] + Vtp))
                    S = S * Yt * Ft
                    population[i, :] = (Gb_Sol + (0.2 * (1 - np.random.rand()) + np.random.rand()) * (U2 * Gb_Sol - population[i, :])) - S

            population[i, :] = np.clip(population[i, :], lb, ub)
            nF = fhd(population[np.newaxis, i, :], func_num)
            if fitness[i] < nF:
                population[i, :] = Xp[i, :]
            else:
                Xp[i, :] = population[i, :]
                fitness[i] = nF.item()
                if nF <= Gb_Fit:
                    Gb_Sol = population[i, :]
                    Gb_Fit = nF.item()

        Conv_curve[t] = Gb_Fit
        t += 1

    return Gb_Fit, Gb_Sol, Conv_curve

# GWO main algorithm---灰狼优化算法---不参与测试
# def GWO(population, pop_size, Tmax, ub, lb, dim, func_num):
#     fitness = fhd(population, func_num)
#     alpha, beta, delta = population[np.argsort(fitness)[:3]]
#     conv_curve = []
#     for t in range(Tmax):
#         a = 2 - t * (2 / Tmax)
#         for i in range(pop_size):
#             r1, r2 = np.random.rand(), np.random.rand()
#             A1 = 2 * a * r1 - a
#             C1 = 2 * r2
#             D_alpha = np.abs(C1 * alpha - population[i])
#             X1 = alpha - A1 * D_alpha
#
#             r1, r2 = np.random.rand(), np.random.rand()
#             A2 = 2 * a * r1 - a
#             C2 = 2 * r2
#             D_beta = np.abs(C2 * beta - population[i])
#             X2 = delta - A2 * D_beta
#
#             r1, r2 = np.random.rand(), np.random.rand()
#             A3 = 2 * a * r1 - a
#             C3 = 2 * r2
#             D_delta = np.abs(C3 * delta - population[i])
#             X3 = delta - A3 * D_delta
#
#             population[i] = (X1 + X2 + X3) / 3
#             population[i] = np.clip(population[i], lb, ub)
#         fitness = fhd(population, func_num)
#         alpha, beta, delta = population[np.argsort(fitness)[:3]]
#         conv_curve.append(fitness[0])
#     return fitness[0], alpha, np.array(conv_curve)
def GWO(pop_size, Tmax, ub, lb, dim, func_num):
    # 常数C1, C2, C3定义
    C1 = 1.5
    C2 = 1.0
    C3 = 1.0

    # 初始化灰狼的位置
    X = initialization(pop_size, dim, ub, lb)
    fitness = fhd(X, func_num)

    # 排序并确定Alpha, Beta, Delta狼
    sorted_indices = np.argsort(fitness)
    alpha_pos = X[sorted_indices[0], :]
    alpha_fit = fitness[sorted_indices[0]]

    beta_pos = X[sorted_indices[1], :]
    beta_fit = fitness[sorted_indices[1]]

    delta_pos = X[sorted_indices[2], :]
    delta_fit = fitness[sorted_indices[2]]

    Conv_curve = np.zeros(Tmax)
    t = 0

    while t < Tmax:
        # 更新灰狼的位置
        a = 2 - t * (2 / Tmax)  # a从2到0线性变化
        for i in range(pop_size):
            # 计算每个狼的当前位置更新
            A1 = 2 * a * np.random.rand(dim) - a
            A2 = 2 * a * np.random.rand(dim) - a
            A3 = 2 * a * np.random.rand(dim) - a

            D_alpha = np.abs(C1 * alpha_pos - X[i, :])
            D_beta = np.abs(C2 * beta_pos - X[i, :])
            D_delta = np.abs(C3 * delta_pos - X[i, :])

            X[i, :] = X[i, :] + (A1 * D_alpha + A2 * D_beta + A3 * D_delta) / 3

            # 更新位置时要进行边界限制
            X[i, :] = np.clip(X[i, :], lb, ub)

        # 计算适应度
        fitness = fhd(X, func_num)

        # 更新Alpha, Beta, Delta狼
        sorted_indices = np.argsort(fitness)
        if fitness[sorted_indices[0]] < alpha_fit:
            alpha_pos = X[sorted_indices[0], :]
            alpha_fit = fitness[sorted_indices[0]]

        if fitness[sorted_indices[1]] < beta_fit:
            beta_pos = X[sorted_indices[1], :]
            beta_fit = fitness[sorted_indices[1]]

        if fitness[sorted_indices[2]] < delta_fit:
            delta_pos = X[sorted_indices[2], :]
            delta_fit = fitness[sorted_indices[2]]

        Conv_curve[t] = alpha_fit  # 保存当前迭代的最优适应度
        t += 1

    return alpha_fit, alpha_pos, Conv_curve
# PSO---粒子群算法---不参与测试


# WOA主算法---鲸鱼优化算法
# def WOA(population, pop_size, Tmax, ub, lb, dim, func_num):
#     fitness = fhd(population, func_num)
#     # 获取最优的鲸鱼位置
#     best_sol = population[np.argmin(fitness)]
#     best_fitness = np.min(fitness)
#     conv_curve = []
#     for t in range(Tmax):
#         a = 2 - t * (2/Tmax)
#         for i in range(pop_size):
#             p = np.random.rand()
#             if p < 0.5:
#                 A = 2 * a * np.random.rand() - a
#                 C = 2 * np.random.rand()
#                 D_leader = np.abs(C * best_sol - population[i])
#                 population[i] = best_sol - A * D_leader
#             else:
#                 l = np.random.uniform(-1, 1)
#                 D_leader = np.abs(best_sol - population[i])
#                 population[i] = D_leader * np.exp(1.5 * l) * np.cos(2 * np.pi * l) + best_sol
#             population[i] = np.clip(population[i], lb, ub)
#         fitness = fhd(population, func_num)
#         current_best_idx = np.argmin(fitness)
#         if fitness[current_best_idx] < best_fitness:
#             best_sol = population[current_best_idx]
#             best_fitness = fitness[current_best_idx]
#         conv_curve.append(best_fitness)
#     return best_fitness, best_sol, np.array(conv_curve)
def WOA(population, pop_size, Tmax, ub, lb, dim, func_num):
    fitness = fhd(population, func_num)

    # 获取最优的鲸鱼位置
    best_idx = np.argmin(fitness)
    best_pos = population[best_idx, :]
    best_fit = fitness[best_idx]

    # 初始化收敛曲线
    Conv_curve = np.zeros(Tmax)

    t = 0
    while t < Tmax:
        a = 2 - t * (2 / Tmax)  # a 从 2 线性变化到 0
        for i in range(pop_size):
            # 选择一个随机值来决定鲸鱼的位置更新方式
            r1 = np.random.rand()
            r2 = np.random.rand()

            A = 2 * a * r1 - a
            C = 2 * r2

            b = 1  # 扩展常数
            l = np.random.uniform(-1, 1)

            p = np.random.rand()

            # 更新鲸鱼位置
            if p < 0.5:
                if np.abs(A) >= 1:
                    D = np.abs(C * best_pos - population[i, :])
                    population[i, :] = best_pos - A * D  # 追踪最优解
                else:
                    D = np.abs(C * best_pos - population[i, :])
                    population[i, :] = best_pos - A * D + b * np.sin(l)  # 利用螺旋更新位置
            else:
                D = np.abs(C * best_pos - population[i, :])
                population[i, :] = best_pos - A * D  # 追踪最优解

            # 边界处理
            population[i, :] = np.clip(population[i, :], lb, ub)

        # 计算适应度
        fitness = fhd(population, func_num)

        # 更新最优解
        best_idx = np.argmin(fitness)
        if fitness[best_idx] < best_fit:
            best_pos = population[best_idx, :]
            best_fit = fitness[best_idx]

        Conv_curve[t] = best_fit  # 记录当前迭代的最优适应度
        t += 1

    return best_fit, best_pos, Conv_curve

# DE算法
def DE(population, pop_size, Tmax, ub, lb, dim, func_num, F = 0.5, CR = 0.9):
    fitness = fhd(population, func_num)
    conv_curve = []
    for t in range(Tmax):
        for i in range(pop_size):
            idxs = [idx for idx in range(pop_size) if idx != i]
            a, b, c = population[np.random.choice(idxs, 3, replace=False)]
            mutant = np.clip(a + F * (b - c), lb, ub)
            cross_points = np.random.rand(dim) < CR
            if not np.any(cross_points):
                cross_points[np.random.randint(0, dim)] = True
            trial = np.where(cross_points, mutant, population[i])
            f_trial = fhd(trial[np.newaxis, :], func_num).item()
            if f_trial < fitness[i]:
                population[i] = trial
                fitness[i] = f_trial
        best_idx = np.argmin(fitness)
        conv_curve.append(fitness[best_idx])
    return fitness[best_idx], population[best_idx], np.array(conv_curve)

# --------------普通遗传算法
def COMGA(population, pop_size, Tmax, ub, lb, dim, func_num, crossover_rate=0.8, mutation_rate=0.2):
    fitness = fhd(population, func_num)
    conv_curve = []
    for t in range(Tmax):
        # Selection
        selected = population[np.argsort(fitness)[:pop_size // 2]]
        # Crossover
        offspring = []
        for i in range(0, len(selected) - 1, 2):
            parent1, parent2 = selected[i], selected[i + 1]
            if np.random.random() < crossover_rate:
                cross_point = np.random.randint(1, dim)
                child1 = np.concatenate((parent1[:cross_point], parent2[cross_point:]))
                child2 = np.concatenate((parent2[:cross_point], parent1[cross_point:]))
                offspring.extend([child1, child2])
            else:offspring.extend([parent1, parent2])
        offspring = np.array(offspring)
        for i in range(len(offspring)):
            if np.random.rand() < mutation_rate:
                mutation_index = np.random.randint(0, dim)
                offspring[i, mutation_index] = np.random.uniform(lb, ub)
        population = np.vstack((selected, offspring))
        fitness = fhd(population, func_num)
        best_idx = np.argmin(fitness)
        conv_curve.append(fitness[best_idx])
    best_fitness = np.min(fitness)
    best_solution = population[np.argmin(fitness)]
    return best_fitness, best_solution, np.array(conv_curve)

# ---------------以下为改进遗传算法算法--------------------
'''版本一'''
def hybrid_GA_CPO(population, pop_size, Tmax, ub, lb, dim, func_num, crossover_rate=0.8, mutation_rate=0.3):
    """
    混合遗传算法与冠豪猪优化算法，用于 CEC2017 测试集。

    参数：
    - population: 初始种群 (numpy array, shape: [pop_size, dim])
    - pop_size: 种群大小
    - Tmax: 最大迭代次数
    - ub: 变量上界
    - lb: 变量下界
    - dim: 问题维度
    - func_num: CEC2017 测试函数编号
    - crossover_rate: 交叉概率，默认 0.8
    - mutation_rate: 变异概率，默认 0.1

    返回：
    - best_fitness: 最佳适应度值
    - best_solution: 最佳解
    - conv_curve: 收敛曲线 (numpy array)
    """

    # 初始化适应度
    fitness = fhd(population, func_num)

    # 初始化全局最优解
    Gb_Fit = np.min(fitness)
    Gb_Sol = population[np.argmin(fitness)].copy()
    conv_curve = []

    # 主循环
    for t in range(Tmax):
        # --- 遗传算法部分 ---
        # 选择：选择适应度最好的前一半个体
        selected = population[np.argsort(fitness)[:pop_size // 2]]

        # 交叉：生成后代
        offspring = []
        for i in range(0, len(selected) - 1, 2):
            parent1, parent2 = selected[i], selected[i + 1]
            if np.random.random() < crossover_rate:
                cross_point = np.random.randint(1, dim)
                child1 = np.concatenate((parent1[:cross_point], parent2[cross_point:]))
                child2 = np.concatenate((parent2[:cross_point], parent1[cross_point:]))
                offspring.extend([child1, child2])
            else:
                offspring.extend([parent1, parent2])
        offspring = np.array(offspring)

        # 变异：对后代应用变异操作
        for i in range(len(offspring)):
            if np.random.rand() < mutation_rate:
                mutation_index = np.random.randint(0, dim)
                offspring[i, mutation_index] = np.random.uniform(lb, ub)

        # 组合形成临时种群
        temporary_population = np.vstack((selected, offspring))

        # 评估临时种群的适应度
        temp_fitness = fhd(temporary_population, func_num)

        # --- 冠豪猪优化部分 ---
        # 对临时种群应用一次 CPO 更新
        for i in range(len(temporary_population)):
            # CPO 更新规则
            U1 = np.random.rand(dim) > 0.5
            rand_index1 = np.random.randint(len(temporary_population))
            rand_index2 = np.random.randint(len(temporary_population))

            if np.random.rand() < 0.5:
                y = (temporary_population[i] + temporary_population[rand_index1]) / 2
                new_candidate = temporary_population[i] + np.random.randn(dim) * np.abs(
                    2 * np.random.rand() * Gb_Sol - y)
            else:
                Yt = 2 * np.random.rand() * (1 - t / Tmax) ** (t / Tmax + 1e-6)
                U2 = np.random.rand(dim) < 0.5
                S = np.random.rand() * U2
                if np.random.rand() < 0.8:
                    St = np.exp(-np.abs(temp_fitness[i] / (np.sum(temp_fitness) + np.finfo(float).eps)))
                    S = S * Yt * St
                    new_candidate = (1 - U1) * temporary_population[i] + U1 * (
                            temporary_population[rand_index1] + St * (
                                temporary_population[rand_index2] - temporary_population[rand_index1]) - S)
                else:
                    Mt = np.exp(-np.abs(temp_fitness[i] / (np.sum(temp_fitness) + np.finfo(float).eps)))
                    Vtp = temporary_population[rand_index1]
                    Ft = np.random.rand(dim) * (Mt * (-temporary_population[i] + Vtp))
                    S = S * Yt * Ft
                    new_candidate = (Gb_Sol + (0.2 * (1 - np.random.rand()) + np.random.rand()) * (
                                U2 * Gb_Sol - temporary_population[i])) - S

            # 边界检查
            new_candidate = np.clip(new_candidate, lb, ub)

            # 评估新候选解的适应度
            new_fitness = fhd(new_candidate[np.newaxis, :], func_num).item()

            # 如果新解更优，则替换原个体
            if new_fitness < temp_fitness[i]:
                temporary_population[i] = new_candidate
                temp_fitness[i] = new_fitness
                # 更新全局最优解
                if new_fitness < Gb_Fit:
                    Gb_Fit = new_fitness
                    Gb_Sol = new_candidate.copy()

        # 更新种群和适应度
        population = temporary_population
        fitness = temp_fitness

        # 记录收敛曲线
        conv_curve.append(Gb_Fit)

    # 返回结果
    best_fitness = Gb_Fit
    best_solution = Gb_Sol
    return best_fitness, best_solution, np.array(conv_curve)

'''版本二'''
"""改进的轮盘赌选择（带惩罚策略）"""
def modified_roulette_selection(population, fitness, current_pop_size, decay_rate=0.03):
    temp_fitness = fitness.copy()  # 使用临时适应度副本
    selected = []

    for _ in range(current_pop_size):
        # Step1: 计算初始概率
        total_fit = np.sum(temp_fitness)
        if total_fit == 0:  # 防止除零错误
            probs = np.ones_like(temp_fitness) / len(temp_fitness)
        else:
            probs = temp_fitness / total_fit

        # Step2: 轮盘赌选择
        selected_idx = np.random.choice(len(population), p=probs)
        selected.append(population[selected_idx])

        # Step3: 应用惩罚（降低被选个体的适应度）
        temp_fitness[selected_idx] *= np.exp(-decay_rate)

        # Step4: 重新归一化概率（避免数值问题）
        total_fit = np.sum(temp_fitness)
        if total_fit != 0:
            temp_fitness = temp_fitness / total_fit

    return np.array(selected)

"""基因块翻转变异"""
def block_reverse_mutation(ind):
    if len(ind) < 2: return ind
    start = np.random.randint(0, len(ind) - 1)
    end = np.random.randint(start + 1, len(ind))
    ind[start:end] = ind[start:end][::-1]
    return ind

"""两点交换变异"""
def two_point_swap_mutation(ind):
    if len(ind) < 2: return ind
    idx1, idx2 = np.random.choice(len(ind), 2, replace=False)
    ind[idx1], ind[idx2] = ind[idx2], ind[idx1]
    return ind

"""相邻基因交换微调"""
def adjacent_swap_mutation(ind):
    if len(ind) < 2: return ind
    idx = np.random.randint(0, len(ind) - 1)
    ind[idx], ind[idx + 1] = ind[idx + 1], ind[idx]
    return ind

"""禁忌爬山算法"""
def taboo_hill_climbing(individual, lb, ub, dim, func_num, taboo_length=5, iterations=5):
    best_ind = individual.copy()
    best_fitness = fhd(best_ind.reshape(1, -1), func_num)[0]
    taboo_list = []
    for _ in range(iterations):
        non_taboo_dims = [d for d in range(dim) if d not in taboo_list]
        if not non_taboo_dims:
            non_taboo_dims = list(range(dim))
        dim_to_change = np.random.choice(non_taboo_dims)
        new_ind = best_ind.copy()
        new_ind[dim_to_change] = np.random.uniform(lb, ub)
        new_fitness = fhd(new_ind.reshape(1, -1), func_num)[0]
        if new_fitness < best_fitness:
            best_ind = new_ind
            best_fitness = new_fitness
            taboo_list.append(dim_to_change)
            if len(taboo_list) > taboo_length:
                taboo_list.pop(0)
    return best_ind

def IMGAOne(population, PN_max, Tmax, ub, lb, dim, func_num, crossover_rate=0.8, mutation_rate=0.2, PN_min = 50, T = 2):
    fitness = fhd(population, func_num) # 种群适应度
    conv_curve = []
    best_individual = population[np.argmin(fitness)] # 记录最优个体
    best_fitness = np.min(fitness) # 记录最优个体适应度
    tracking_best_counter = 0

    for t in range(Tmax):
        # 动态调整种群数量
        G_cycle = Tmax // T
        g_cycle = t % G_cycle
        current_pop_size = int(PN_max - (PN_max - PN_min) * (g_cycle / G_cycle)) # 下一代种群数量

        # 改进的轮盘赌选择（带惩罚策略）
        selected = modified_roulette_selection(population, fitness, current_pop_size)

        # PMX交叉（保持原实现）
        offspring = [] # 保存当前代种群
        for i in range(0, current_pop_size - 1, 2):
            parent1, parent2 = selected[i], selected[i + 1]
            if np.random.rand() < crossover_rate:
                cross_point1 = np.random.randint(0, dim)
                cross_point2 = np.random.randint(cross_point1, dim)
                child1 = np.concatenate(
                    [parent1[:cross_point1], parent2[cross_point1:cross_point2], parent1[cross_point2:]])
                child2 = np.concatenate(
                    [parent2[:cross_point1], parent1[cross_point1:cross_point2], parent2[cross_point2:]])
                offspring.extend([child1, child2])
            else:
                offspring.extend([parent1.copy(), parent2.copy()])
        offspring = np.array(offspring)

        # 分层变异策略
        if len(offspring) > 0:
            offspring_fitness = fhd(offspring, func_num)
            sorted_indices = np.argsort(offspring_fitness)
            ranks = np.argsort(sorted_indices)
            for i in range(len(offspring)):
                percentile = ranks[i] / len(offspring)
                if np.random.rand() < mutation_rate:
                    # 根据排名选择变异策略
                    if percentile >= 0.6:  # 后40%
                        offspring[i] = block_reverse_mutation(offspring[i].copy())
                    elif percentile >= 0.4:  # 40%-60%
                        offspring[i] = two_point_swap_mutation(offspring[i].copy())
                    else:  # 前40%
                        offspring[i] = adjacent_swap_mutation(offspring[i].copy())

        current_fitness = fhd(offspring, func_num) # 获取当前种群适应度值
        current_best_idx = np.argmin(current_fitness) # 获取最优个体索引
        current_best_individual = offspring[current_best_idx] # 获取最优个体
        if current_fitness[current_best_idx] == best_fitness:
            tracking_best_counter += 1
            if tracking_best_counter >= 3:  # 连续3代未改进进行局部破坏重构
                reset_dims = np.random.choice(dim, size=int(dim * 0.3), replace=False)
                current_best_individual[reset_dims] = np.random.uniform(lb, ub, size=len(reset_dims))
                offspring[current_best_idx] = current_best_individual.copy()
                current_fitness = fhd(offspring, func_num)  # 获取当前种群适应度值
                current_best_individual = offspring[np.argmin(current_fitness)] # 更新最优个体
                if current_fitness[np.argmin(current_fitness)] <= best_fitness:
                    best_individual = current_best_individual
                    best_fitness = current_fitness[np.argmin(current_fitness)]
                tracking_best_counter = 0

        for i in range(len(offspring)):
            offspring[i] = taboo_hill_climbing(offspring[i], lb, ub, dim, func_num)
        current_fitness = fhd(offspring, func_num)  # 获取当前种群适应度值
        current_best_idx = np.argmin(current_fitness)  # 获取最优个体索引
        current_best_individual = offspring[current_best_idx]  # 获取最优个体
        if current_fitness[current_best_idx] <= best_fitness:
            best_individual = current_best_individual
            best_fitness = current_fitness[current_best_idx]
        # 合并种群并保留最优
        population = np.vstack((selected, offspring, best_individual.reshape(1, -1)))
        fitness = fhd(population, func_num)
        sorted_indices = np.argsort(fitness)[:current_pop_size]
        population = population[sorted_indices]
        fitness = fitness[sorted_indices]
        conv_curve.append(best_fitness)

    return best_fitness, best_individual, np.array(conv_curve)